package com.woniuxy.rag;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.FileSystemResource;
import org.springframework.stereotype.Component;
import org.springframework.web.bind.annotation.GetMapping;

import java.io.File;
import java.nio.file.Paths;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;

/**
 * @Author: 马宇航
 * @Description: 文件向量化处理工具类
 * @DateTime: 25/10/31/星期五 17:17
 * @Component: 成都蜗牛学苑
 **/
@Component
public class FileProcessVectorTools {
    @Autowired
    PgVectorStore vectorStore;
    @Autowired
    ChatClient chatClient;
    @Value("${woniu.filePathTmp}")
    String path;

    @Tool(description = "文件向量化处理工具类，帮我向量化一下文件")
    public String vector(@ToolParam(description = "只需要文件名，不要路径名") String ...fileNames){
        //1.遍历所有文件，根据文件名判断执行不同的文件处理方案（图片也行）
        for (String fileName : fileNames) {
            //2.读取文件，动态分隔符File.separator
            FileSystemResource fileSystemResource = new FileSystemResource(Paths.get(path + File.separator + fileName));
            //3.判断文件后缀
            String filename = Objects.requireNonNull(fileSystemResource.getFilename());
            //4.判空+后缀
            if (filename.endsWith(".pdf")) {
                //5. pdf阅读器
                PagePdfDocumentReader pagePdfDocumentReader = new PagePdfDocumentReader(fileSystemResource,
                        PdfDocumentReaderConfig.builder()
                                .withPageExtractedTextFormatter(ExtractedTextFormatter.defaults())
                                .withPagesPerDocument(5)//每页做一个向量
                                .build());
                //6. read后，就可以拿到文本内容
                List<Document> documents = pagePdfDocumentReader.read();
                vectorStore.add(documents);
                //tika的word处理
            }else if(filename.endsWith(".doc")||filename.endsWith(".docx")){
                //tika的word处理
                TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(fileSystemResource, ExtractedTextFormatter.builder()
                        .withNumberOfBottomTextLinesToDelete(1) //页尾
                        .withNumberOfTopTextLinesToDelete(1) //页头
                        .build()
                );
                List<Document> documents = tikaDocumentReader.read();
                vectorStore.add(documents);
                //处理md
            }
        }
        return "向量新增成功！";
    }
}
